#include <ros/ros.h>
#include <rosbag/recorder.h>
#include <std_msgs/Bool.h>
#include <mavros_msgs/State.h>
#include <yaml-cpp/yaml.h>
#include <fstream>
#include <iostream>
#include <sys/stat.h>
#include <ctime>
#include <thread>

// 获取当前时间字符串，用于文件命名
std::string getCurrentTimeString() {
    time_t now = time(nullptr);
    char buf[100];
    strftime(buf, sizeof(buf), "%Y-%m-%d-%H-%M-%S", localtime(&now));
    return std::string(buf);
}

// 从 YAML 配置文件加载话题列表
std::vector<std::string> loadTopicsFromYAML(const std::string &yaml_file) {
    YAML::Node config = YAML::LoadFile(yaml_file);
    std::vector<std::string> topics;
    for (YAML::const_iterator it = config["topics"].begin(); it != config["topics"].end(); ++it) {
        topics.push_back(it->as<std::string>());
    }
    return topics;
}

bool startRecording = false;
rosbag::Recorder* recorder = nullptr;
std::string bag_directory;
std::thread recording_thread;
std::shared_ptr<rosbag::Recorder> recorder_ptr;

void recordingFunction(rosbag::RecorderOptions options) {
    recorder_ptr.reset( new rosbag::Recorder(options));
    recorder_ptr->run();
}

// 回调函数，用于启动/停止录制
void toggleRecordingCallback(const mavros_msgs::State::ConstPtr& msg) {
    if (msg->armed) {
        if (!startRecording) {
            startRecording = true;
            std::string bag_file = bag_directory + "/" + "record_" + getCurrentTimeString() + ".bag";

            std::string config_file = "./config/config.yaml";
            char *config_env = getenv("ROS_RECORDER_CONFIG_PATH");
            if (config_env != nullptr) {
                config_file = config_env;
            }
            std::vector<std::string> topics = loadTopicsFromYAML(config_file);

            rosbag::RecorderOptions options;
            options.record_all = false;
            options.regex = false;
            options.topics = topics;
            options.prefix = bag_file;
            options.append_date = false;

            recording_thread = std::thread(recordingFunction, options);
            ROS_INFO("Recording started: %s", bag_file.c_str());
        }
    } else {
        if (startRecording) {
            startRecording = false;
            if (recorder_ptr)
            {
                recorder_ptr->stop();
            }
            ROS_INFO("Recording stopped.");
        }
    }
}

int main(int argc, char **argv) {
    std::string drone_id_str = "uav";
    char *drone_id = getenv("DRONE_ID"); // kml的管理直接放在传输文件夹中，不转移
    if (drone_id != nullptr) {
        drone_id_str = drone_id_str + std::string(drone_id);
    }
    
    ros::init(argc, argv, (drone_id_str + "_rosbag_recorder_node").c_str());
    ros::NodeHandle nh(drone_id_str);
    // 创建保存 bag 文件的目录，如果不存在则创建
    // 创建保存 bag 文件的目录，如果不存在则创建
    bag_directory = "./bag";
    char *bag_dir_env = getenv("ROS_RECORDER_BAG_PATH");
    if (bag_dir_env != nullptr) {
	    bag_directory = bag_dir_env;
    }

    mkdir(bag_directory.c_str(), 0755);
    bag_directory += "/" + getCurrentTimeString();
    mkdir(bag_directory.c_str(), 0755);

    // 订阅 toggle_recording 话题
    ros::Subscriber sub = nh.subscribe("mavros/state", 10, toggleRecordingCallback);

    // 循环等待外部信号
    ros::spin();

    if (recorder_ptr)
    {
        recorder_ptr->stop();
    }
    ROS_INFO("Recording stopped.");

    if (recording_thread.joinable())
    {
        recording_thread.join();
    }

    return 0;
}
